#!/usr/bin/python

'''
Goal: use a list values (scores) to fit into negative binomail distribution,
Then calculate a threshould based on user defined pvalue;
If the data cannot be fit into a NB distribution, use Empirical distribution.
** Choose only scores less than 1000 (to eliminate the repetative regions or PCR duplicate?)
'''


from __future__ import division
import re, os, sys, shutil
from math import *
from string import *
from optparse import OptionParser
import operator
import numpy as np
from scipy.stats import nbinom
from scipy.special import psi
from scipy.optimize import brentq


def CalculateThreshold(infile, pvalue):
    scores = [];
    scorefile = open(infile, 'r');
    for lines in scorefile:
        lines = lines.strip().split();
        if int(lines[0]) < 1000:         # this is to eliminate the repeat regions(?)
            scores.append(int(lines[0]));
    try:
        (size, prop) = fit_nbinom(np.array(scores));
        threshold = nbinom.ppf(1-pvalue, size, prop);
        print threshold;
    except:
        threshold = np.percentile(np.array(scores), (1-pvalue)*100);
        print threshold;


def fit_nbinom(k):
    N = len(k)
    n = brentq(lambda r: sum(psi(k + r)) - N*psi(r) + N*log(r/(r + sum(k/N))), np.finfo(np.float).eps, np.max(k))
    p = n/(n + sum(k/N))
    return n, p


def main(argv):
    parser = OptionParser()
    parser.add_option("-i", "--inputfile", action="store", type="string",
                      dest="infile", help="input file with counts", metavar="<file>")
    parser.add_option("-p", "--p-value", action="store", type="float",
                      dest="pvalue", help="p-value for searching potential enriched window, default: 0.02",
                      metavar="<float>", default=0.02)

    (opt, args) = parser.parse_args(argv)
    if len(argv) < 2:
        parser.print_help()
        sys.exit(1)

    CalculateThreshold(opt.infile, opt.pvalue);


if __name__ == "__main__":
    main(sys.argv)
